from tensorflow.keras.models import load_model
from tensorflow.keras.utils import CustomObjectScope
from tensorflow.keras.initializers import glorot_uniform
import os

def loadModel(modelPath):
    with CustomObjectScope({'GlorotUniform': glorot_uniform()}):
        model = load_model(modelPath)
    return model

def predict_function(modelPath):
    model = None
    with CustomObjectScope({'GlorotUniform': glorot_uniform()}):
        model = load_model(modelPath)
    def predict(input):
        return model.predict(input, 256)
    return predict

def enableGPU(turnOn=True):
    if turnOn:
        os.environ["CUDA_VISIBLE_DEVICES"] = "0"
    else:
        os.environ["CUDA_VISIBLE_DEVICES"] = "-1"